import os
import sys
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
                             QLabel, QLineEdit, QPushButton, QListWidget, QGraphicsView,
                             QGraphicsScene, QGraphicsPixmapItem, QTextEdit, QSpinBox,
                             QFileDialog, QComboBox, QDialog, QDialogButtonBox, QInputDialog,
                             QGroupBox, QFormLayout, QMessageBox, QColorDialog, QSlider, QGraphicsRectItem, QGraphicsPathItem, QGraphicsPolygonItem, QGraphicsTextItem)
from PyQt5.QtGui import (QPixmap, QImage, QPainter, QPen, QColor, QCursor, QFont,
                         QPolygonF, QBrush, QPainterPath)
from PyQt5.QtCore import Qt, QRectF, QPointF, QEvent
from PIL import Image, ImageDraw, ImageFont
import re

import pandas as pd


def resource_path(relative_path):
    """获取资源的绝对路径"""
    if hasattr(sys, '_MEIPASS'):
        base_path = sys._MEIPASS
    else:
        base_path = os.path.abspath(".")
    return os.path.join(base_path, relative_path)

# 读取excel文件
# _current_dir = os.path.dirname(os.path.abspath(__file__))
# _path = os.path.join(_current_dir, 'static', 'config.xlsx')
df = pd.read_excel(resource_path("static/config.xlsx"))  # 替换为你的文件路径和文件名


class ColorWidthDialog(QDialog):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowTitle("设置颜色和线宽")

        layout = QVBoxLayout()

        # 颜色选择
        self.color_label = QLabel("选择颜色:")
        self.color_btn = QPushButton("选择颜色")
        self.color_btn.clicked.connect(self.choose_color)
        self.color_preview = QLabel()
        self.color_preview.setFixedSize(50, 20)
        self.current_color = QColor(255, 0, 0)  # 默认红色
        self.update_color_preview()

        # 线宽选择
        self.width_label = QLabel("线宽:")
        self.width_slider = QSlider(Qt.Horizontal)
        self.width_slider.setRange(5, 15)
        self.width_slider.setValue(15)
        self.width_value = QLabel("15")

        # 表单布局
        form = QFormLayout()
        form.addRow(self.color_label, self.color_btn)
        form.addRow(self.color_preview)
        form.addRow(self.width_label)
        form.addRow(self.width_slider)
        form.addRow(self.width_value)

        # 按钮
        buttons = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel)
        buttons.button(QDialogButtonBox.StandardButton.Ok).setText("确定")
        buttons.button(QDialogButtonBox.StandardButton.Cancel).setText("取消")
        buttons.accepted.connect(self.accept)
        buttons.rejected.connect(self.reject)

        # 连接信号
        self.width_slider.valueChanged.connect(self.update_width_value)

        layout.addLayout(form)
        layout.addWidget(buttons)
        self.setLayout(layout)

    def choose_color(self):
        color = QColorDialog.getColor(self.current_color, self, "选择颜色")
        if color.isValid():
            self.current_color = color
            self.update_color_preview()

    def update_color_preview(self):
        self.color_preview.setStyleSheet(f"background-color: {self.current_color.name()};")

    def update_width_value(self, value):
        self.width_value.setText(str(value))

    def get_settings(self):
        return self.current_color, self.width_slider.value()


class AnnotationDialog(QDialog):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowTitle("保存标注信息")

        layout = QVBoxLayout()

        # 错误位置
        self.position_combo = QComboBox()
        position_column_data = df['出错位置']
        position_data_array = position_column_data.dropna().values.tolist()
        self.position_combo.addItems(position_data_array)
        self.position_combo.currentTextChanged.connect(self.position_combo_currentTextChanged)

        # 错误类型
        self.type_combo = QComboBox()
        typpe_column_data = df['出错类型']
        type_data_array = typpe_column_data.dropna().values.tolist()
        filtered_strings = [s for s in type_data_array if s.find(self.position_combo.currentText()) != -1]
        self.type_combo.addItems(filtered_strings)
        self.type_combo.currentTextChanged.connect(self.type_combo_currentTextChanged)

        # 危险等级
        self.level_combo = QComboBox()
        self.level_combo.addItems(["一般", "严重", "危急"])
        self.level_combo.currentTextChanged.connect(self.level_combo_currentTextChanged)

        # 自定义描述
        self.desc_input = QLineEdit()
        self.desc_input.setPlaceholderText("自定义重命名")
        self.desc_input.setText(self.type_combo.currentText() + '_' + self.level_combo.currentText())

        # 表单布局
        form = QFormLayout()
        form.addRow("出错位置:", self.position_combo)
        form.addRow("出错类型:", self.type_combo)
        form.addRow("危险等级:", self.level_combo)
        form.addRow("自定义重命名:", self.desc_input)

        # 按钮

        buttons = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel)
        buttons.button(QDialogButtonBox.StandardButton.Ok).setText("确定")
        buttons.button(QDialogButtonBox.StandardButton.Cancel).setText("取消")
        buttons.accepted.connect(self.accept)
        buttons.rejected.connect(self.reject)

        layout.addLayout(form)
        layout.addWidget(buttons)
        self.setLayout(layout)

    def position_combo_currentTextChanged(self, label):
        typpe_column_data = df['出错类型']
        type_data_array = typpe_column_data.dropna().values.tolist()
        filtered_strings = [s for s in type_data_array if s.find(label) != -1]
        self.type_combo.clear()
        self.type_combo.addItems(filtered_strings)
        self.type_combo.repaint()
        self.desc_input.setText(self.type_combo.currentText() + '_' + self.level_combo.currentText())

    def type_combo_currentTextChanged(self, label):
        self.desc_input.setText(self.type_combo.currentText() + '_' + self.level_combo.currentText())

    def level_combo_currentTextChanged(self, label):
        self.desc_input.setText(self.type_combo.currentText() + '_' + self.level_combo.currentText())

    def get_info(self):
        return {
            'position': self.position_combo.currentText(),
            'type': self.type_combo.currentText(),
            'level': self.level_combo.currentText(),
            'desc': self.desc_input.text()
        }


class ImageViewer(QGraphicsView):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.scene = QGraphicsScene(self)
        self.setScene(self.scene)

        self.pixmap_item = None
        self.current_scale = 1.0
        self.dragging = False
        self.last_pos = QPointF()

        # 标注相关
        self.drawing = False
        self.current_shape = None
        self.shapes = []
        self.current_annotations = []
        self.annotation_mode = None
        self.start_pos = None
        self.is_annotating = False
        self.polygon_points = []

        # 默认样式
        self.current_color = QColor(255, 0, 0)
        self.current_width = 15

        # 设置鼠标跟踪
        self.setMouseTracking(True)
        self.setRenderHint(QPainter.Antialiasing)

    def set_image(self, image_path):
        self.scene.clear()
        self.shapes.clear()
        self.current_annotations.clear()
        self.polygon_points.clear()

        pixmap = QPixmap(image_path)
        if pixmap.isNull():
            return False

        self.pixmap_item = QGraphicsPixmapItem(pixmap)
        self.scene.addItem(self.pixmap_item)
        self.setSceneRect(QRectF(pixmap.rect()))

        self.current_scale = 1.0
        self.fitInView(self.sceneRect(), Qt.KeepAspectRatio)
        return True

    def wheelEvent(self, event):
        if not self.is_annotating:  # 仅在非标注模式下允许缩放
            zoom_factor = 1.2
            if event.angleDelta().y() > 0:
                self.scale(zoom_factor, zoom_factor)
                self.current_scale *= zoom_factor
            else:
                self.scale(1.0 / zoom_factor, 1.0 / zoom_factor)
                self.current_scale /= zoom_factor

    def mousePressEvent(self, event):
        if event.button() == Qt.LeftButton:
            if self.is_annotating and self.annotation_mode:
                # 开始绘制标注
                self.drawing = True
                self.start_pos = self.mapToScene(event.pos())
                print('self.start_pos', self.start_pos)
                if self.annotation_mode in ['rectangle', 'circle']:
                    self.current_shape = self.scene.addRect(QRectF(self.start_pos, self.start_pos))
                    self.current_shape.setPen(QPen(self.current_color, self.current_width))
                elif self.annotation_mode == 'text':
                    # 文本标注将在鼠标释放时处理
                    pass
                elif self.annotation_mode == 'polygon':
                    if not self.current_shape:
                        self.polygon_points = [self.start_pos]
                        self.current_shape = self.scene.addPolygon(QPolygonF(self.polygon_points),
                                                                   QPen(self.current_color, self.current_width))
                    else:
                        self.polygon_points.append(self.start_pos)
                        self.current_shape.setPolygon(QPolygonF(self.polygon_points))
            elif not self.is_annotating:
                # 开始拖动
                self.dragging = True
                self.last_pos = event.pos()
                self.setCursor(Qt.ClosedHandCursor)
        elif event.button() == Qt.RightButton and self.annotation_mode == 'polygon' and self.is_annotating:
            # 完成多边形绘制
            if self.current_shape and len(self.polygon_points) >= 3:
                self.shapes.append(self.current_shape)
                self.current_shape = None
                self.polygon_points = []
                self.drawing = False

        super().mousePressEvent(event)

    def mouseMoveEvent(self, event):
        if self.drawing and self.is_annotating and self.annotation_mode in ['rectangle', 'circle']:
            # 更新当前形状
            end_pos = self.mapToScene(event.pos())

            rect = QRectF(self.start_pos, end_pos).normalized()

            if self.annotation_mode == 'rectangle':
                self.current_shape.setRect(rect)
            elif self.annotation_mode == 'circle':

                # 以起点为中心，计算半径
                radius = QPointF(end_pos - self.start_pos).manhattanLength()
                # self.current_item.setRect(QRectF(
                #     self.start_pos.x() - radius, self.start_pos.y() - radius,
                #     radius * 2, radius * 2
                # ))

                path = self.create_ellipse_path(rect, radius)
                if self.current_shape:
                    self.scene.removeItem(self.current_shape)
                self.current_shape = self.scene.addPath(path, QPen(self.current_color, self.current_width))

        elif self.dragging and not self.is_annotating:
            # 拖动视图
            delta = event.pos() - self.last_pos
            self.last_pos = event.pos()
            self.horizontalScrollBar().setValue(self.horizontalScrollBar().value() - delta.x())
            self.verticalScrollBar().setValue(self.verticalScrollBar().value() - delta.y())

        super().mouseMoveEvent(event)

    def mouseReleaseEvent(self, event):
        if event.button() == Qt.LeftButton and self.drawing and self.is_annotating:
            if self.annotation_mode == 'rectangle':
                end_pos = self.mapToScene(event.pos())
                rect = QRectF(self.start_pos, end_pos).normalized()
                if rect.width() > 5 and rect.height() > 5:  # 最小尺寸
                    self.shapes.append(self.current_shape)
                else:
                    self.scene.removeItem(self.current_shape)

            elif self.annotation_mode == 'circle':
                end_pos = self.mapToScene(event.pos())
                rect = QRectF(self.start_pos, end_pos).normalized()
                if rect.width() > 5 and rect.height() > 5:  # 最小尺寸
                    # 以起点为中心，计算半径
                    radius = QPointF(end_pos - self.start_pos).manhattanLength()
                    path = self.create_ellipse_path(rect, radius)
                    self.scene.removeItem(self.current_shape)
                    self.current_shape = self.scene.addPath(path, QPen(self.current_color, self.current_width))
                    self.shapes.append(self.current_shape)
                else:
                    self.scene.removeItem(self.current_shape)

            elif self.annotation_mode == 'text':
                pos = self.mapToScene(event.pos())
                text, ok = QInputDialog.getText(self, "输入文本", "请输入标注文本:")
                if ok and text:
                    text_item = self.scene.addText(text)
                    text_item.setPos(pos)
                    text_item.setDefaultTextColor(self.current_color)
                    font = QFont()
                    font.setPointSize(12)
                    text_item.setFont(font)
                    self.shapes.append(text_item)

            self.drawing = False
            self.current_shape = None

        elif event.button() == Qt.LeftButton and self.dragging:
            self.dragging = False
            self.setCursor(Qt.ArrowCursor)

        super().mouseReleaseEvent(event)

    def create_ellipse_path(self, rect, radius):
        path = QPainterPath()
        # path.addEllipse(rect)
        # path.addEllipse(QRectF(rect.x(), rect.y(), rect.height(), rect.height()))

        path.addEllipse(QRectF(self.start_pos.x() - radius, self.start_pos.y() - radius, radius * 2, radius * 2))
        return path

    def clear_annotations(self):
        for shape in self.shapes:
            self.scene.removeItem(shape)
        self.shapes.clear()
        self.polygon_points.clear()

    def contains_hashtag_number(self, filename):
        """检查字符串中是否包含#后跟数字的模式（如#009）"""
        return bool(re.search(r'#\d+', filename))

    def save_annotated_image(self, image_path, info, folder_path):
        if not self.pixmap_item:
            return None, None

        # 获取原始图片
        pil_img = Image.open(image_path).convert("RGB")
        img_width, img_height = pil_img.size

        # 绘制标注
        draw = ImageDraw.Draw(pil_img)
        try:
            font = ImageFont.truetype("arial.ttf", 20)
        except:
            font = ImageFont.load_default()


        for shape in self.shapes:
            if isinstance(shape, QGraphicsRectItem):
                rect = shape.rect()
                x1, y1 = rect.x(), rect.y()
                x2, y2 = x1 + rect.width(), y1 + rect.height()
                color = shape.pen().color()
                width = shape.pen().width()
                draw.rectangle([x1, y1, x2, y2],
                               outline=(color.red(), color.green(), color.blue()),
                               width=width)

            elif isinstance(shape, QGraphicsPathItem):
                path = shape.path()
                rect = path.boundingRect()
                x1, y1 = rect.x(), rect.y()
                x2, y2 = x1 + rect.width(), y1 + rect.height()
                color = shape.pen().color()
                width = shape.pen().width()
                draw.ellipse([x1, y1, x2, y2],
                             outline=(color.red(), color.green(), color.blue()),
                             width=width)

            elif isinstance(shape, QGraphicsPolygonItem):
                polygon = shape.polygon()
                points = [(p.x(), p.y()) for p in polygon]
                color = shape.pen().color()
                width = shape.pen().width()
                if len(points) >= 3:
                    draw.polygon(points,
                                 outline=(color.red(), color.green(), color.blue()),
                                 width=width)

            elif isinstance(shape, QGraphicsTextItem):
                text = shape.toPlainText()
                pos = shape.pos()
                color = shape.defaultTextColor()
                draw.text((pos.x(), pos.y()), text,
                          fill=(color.red(), color.green(), color.blue()),
                          font=font)


        _parent_dir = os.path.dirname(folder_path)
        _qx_full_path = os.path.join(_parent_dir, '缺陷图片')

        # 检查目录是否存在
        if not os.path.exists(_qx_full_path):
            # 创建目录
            os.makedirs(_qx_full_path, exist_ok=True)

        #     判断缺陷图片下面是否有一般，严重，危急文件夹
        _yb_full_path = os.path.join(_qx_full_path, '一般')
        _yz_full_path = os.path.join(_qx_full_path, '严重')
        _wj_full_path = os.path.join(_qx_full_path, '危急')
        if not os.path.exists(_yb_full_path):
            os.makedirs(_yb_full_path, exist_ok=True)
        if not os.path.exists(_yz_full_path):
            os.makedirs(_yz_full_path, exist_ok=True)
        if not os.path.exists(_wj_full_path):
            os.makedirs(_wj_full_path, exist_ok=True)

        # 图片名称
        _pic_name = os.path.basename(image_path)
        _ress = _pic_name.split('.')
        if self.contains_hashtag_number(_ress[0]) is False:
            _c_path = os.path.basename(os.path.dirname(image_path))
            _ress[0] = _c_path

        annotated_path = ''
        # 保存标注后的完整图片
        if info['level'] == '一般':
            annotated_path = os.path.join(_yb_full_path, _ress[0] + '_' + info['desc'] + '.' + _ress[1])
        if info['level'] == '严重':
            annotated_path = os.path.join(_yz_full_path, _ress[0] + '_' + info['desc'] + '.' + _ress[1])
        if info['level'] == '危急':
            annotated_path = os.path.join(_wj_full_path, _ress[0] + '_' + info['desc'] + '.' + _ress[1])

        pil_img.save(annotated_path)

        # 6. 裁剪标注区域
        if self.shapes:
            # 计算所有标注的边界框（使用图片坐标）
            min_x, min_y = float('inf'), float('inf')
            max_x, max_y = float('-inf'), float('-inf')

            for shape in self.shapes:
                if isinstance(shape, (QGraphicsRectItem, QGraphicsPathItem, QGraphicsPolygonItem)):
                    rect = shape.boundingRect()
                    # top_left = self.mapFromScene(rect.topLeft())
                    # bottom_right = self.mapFromScene(rect.bottomRight())
                    top_left = rect.topLeft()
                    bottom_right = rect.bottomRight()

                    min_x = min(min_x, top_left.x())
                    min_y = min(min_y, top_left.y())
                    max_x = max(max_x, bottom_right.x())
                    max_y = max(max_y, bottom_right.y())
                elif isinstance(shape, QGraphicsTextItem):
                    pos = shape.pos()
                    view_pos = self.mapFromScene(pos)
                    rect = shape.boundingRect()

                    min_x = min(min_x, view_pos.x())
                    min_y = min(min_y, view_pos.y())
                    max_x = max(max_x, view_pos.x() + rect.width())
                    max_y = max(max_y, view_pos.y() + rect.height())

            # 计算中心点和裁剪区域
            center_x = (min_x + max_x) / 2
            center_y = (min_y + max_y) / 2
            region_width = max_x - min_x
            region_height = max_y - min_y

            # 添加边距（至少50像素）
            margin = max(50, region_width * 0.2, region_height * 0.2)
            crop_size = max(region_width, region_height) + 2 * margin

            # 计算裁剪区域（确保在图片范围内）
            crop_min_x = max(0, center_x - crop_size / 2)
            crop_min_y = max(0, center_y - crop_size / 2)
            crop_max_x = min(img_width, center_x + crop_size / 2)
            crop_max_y = min(img_height, center_y + crop_size / 2)

            # 调整确保有效区域
            if crop_max_x <= crop_min_x:
                crop_max_x = crop_min_x + 100  # 最小宽度100像素

            if crop_max_y <= crop_min_y:
                crop_max_y = crop_min_y + 100  # 最小高度100像素

            # 裁剪并保存
            crop_box = (int(crop_min_x), int(crop_min_y),
                        int(crop_max_x), int(crop_max_y))

            try:
                cropped_img = pil_img.crop(crop_box)

                cropped_path = ''
                # 根据危险等级创建文件夹
                if info['level'] == '一般':
                    cropped_path = os.path.join(_yb_full_path, _ress[0] + '_' + info['desc'] + '(2)' + '.' + _ress[1])
                if info['level'] == '严重':
                    cropped_path = os.path.join(_yz_full_path, _ress[0] + '_' + info['desc'] + '(2)' + '.' + _ress[1])
                if info['level'] == '危急':
                    cropped_path = os.path.join(_wj_full_path, _ress[0] + '_' + info['desc'] + '(2)' + '.' + _ress[1])

                cropped_img.save(cropped_path)

                return annotated_path, cropped_path
            except Exception as e:
                self.output_text.append(f"裁剪图片时出错: {str(e)}")
                return annotated_path, None

        return annotated_path, None

    def set_annotating(self, annotating):
        self.is_annotating = annotating
        if annotating:
            self.setDragMode(QGraphicsView.NoDrag)
        else:
            self.setDragMode(QGraphicsView.ScrollHandDrag)

    def set_pen_style(self, color, width):
        self.current_color = color
        self.current_width = width


class ImageAnnotationWidget(QWidget):
    def __init__(self):
        super().__init__()
        self.initUI()
        self.current_image_index = -1
        self.image_files = []
        self.has_annotator_files = []

    def initUI(self):
        layout = QHBoxLayout()

        # 左侧面板 - 图片列表和设置 (占20%)
        left_panel = QVBoxLayout()
        left_panel.setContentsMargins(5, 5, 5, 5)

        # 文件夹选择
        folder_group = QGroupBox("图片选择")
        folder_layout = QVBoxLayout()

        self.folder_label = QLabel("选择图片文件夹:")
        self.folder_path = QLineEdit()
        self.folder_btn = QPushButton("浏览...")
        self.folder_btn.clicked.connect(self.select_folder)

        folder_top_layout = QHBoxLayout()
        folder_top_layout.addWidget(self.folder_label)
        folder_top_layout.addWidget(self.folder_path)
        folder_top_layout.addWidget(self.folder_btn)

        # 起始图片设置
        self.start_label = QLabel("从第几张图片开始加载:")
        self.start_spin = QSpinBox()
        self.start_spin.setMinimum(1)
        self.start_spin.setValue(1)

        start_layout = QHBoxLayout()
        start_layout.addWidget(self.start_label)
        start_layout.addWidget(self.start_spin)

        # 图片列表
        self.image_list = QListWidget()
        self.image_list.itemClicked.connect(self.on_image_selected)

        # 添加上一张/下一张按钮
        nav_buttons = QHBoxLayout()
        self.prev_btn = QPushButton("上一张")
        self.prev_btn.clicked.connect(self.loadPreviousImage)
        self.next_btn = QPushButton("下一张")
        self.next_btn.clicked.connect(self.loadNextImage)
        nav_buttons.addWidget(self.prev_btn)
        nav_buttons.addWidget(self.next_btn)



        # 添加到文件夹组
        folder_layout.addLayout(folder_top_layout)
        folder_layout.addLayout(start_layout)
        folder_layout.addWidget(self.image_list)
        folder_layout.addLayout(nav_buttons)
        folder_group.setLayout(folder_layout)


        # 添加到左侧面板
        left_panel.addWidget(folder_group)

        # 中间面板 - 图片显示 (占60%)
        self.viewer = ImageViewer()

        # 右侧面板 - 标注工具和输出 (占20%)
        right_panel = QVBoxLayout()
        right_panel.setContentsMargins(5, 5, 5, 5)

        # 标注工具
        annotate_group = QGroupBox("标注工具")
        annotate_layout = QVBoxLayout()

        # 样式设置按钮
        self.style_btn = QPushButton("设置颜色和线宽")
        self.style_btn.clicked.connect(self.set_style)

        # 标注类型按钮
        self.rect_btn = QPushButton("矩形标注")
        self.rect_btn.clicked.connect(lambda: self.set_annotation_mode('rectangle'))

        self.circle_btn = QPushButton("圆形标注")
        self.circle_btn.clicked.connect(lambda: self.set_annotation_mode('circle'))

        self.text_btn = QPushButton("文本标注")
        self.text_btn.clicked.connect(lambda: self.set_annotation_mode('text'))

        self.polygon_btn = QPushButton("多边形标注")
        self.polygon_btn.clicked.connect(lambda: self.set_annotation_mode('polygon'))

        # 操作按钮
        self.start_annotate_btn = QPushButton("开始标绘")
        self.start_annotate_btn.clicked.connect(self.toggle_annotation)

        self.clear_btn = QPushButton("清除标绘")
        self.clear_btn.clicked.connect(self.clear_annotations)

        self.save_btn = QPushButton("保存标绘")
        self.save_btn.clicked.connect(self.save_annotations)

        # 设置按钮样式
        for btn in [self.style_btn, self.rect_btn, self.circle_btn, self.text_btn,
                    self.polygon_btn, self.start_annotate_btn, self.clear_btn, self.save_btn]:
            btn.setStyleSheet("QPushButton { padding: 5px; }")

        # 添加到布局
        annotate_layout.addWidget(self.style_btn)
        annotate_layout.addWidget(self.rect_btn)
        annotate_layout.addWidget(self.circle_btn)
        # annotate_layout.addWidget(self.text_btn)
        # annotate_layout.addWidget(self.polygon_btn)
        annotate_layout.addWidget(self.start_annotate_btn)
        annotate_layout.addWidget(self.clear_btn)
        annotate_layout.addWidget(self.save_btn)
        annotate_layout.addStretch()
        annotate_group.setLayout(annotate_layout)

        # 输出区域
        output_group = QGroupBox("输出信息")
        output_layout = QVBoxLayout()
        self.output_text = QTextEdit()
        self.output_text.setReadOnly(True)
        output_layout.addWidget(self.output_text)
        output_group.setLayout(output_layout)

        # 添加到右侧面板
        right_panel.addWidget(annotate_group)
        right_panel.addWidget(output_group)

        # 添加到主布局
        layout.addLayout(left_panel, 2)  # 20%宽度
        layout.addWidget(self.viewer, 6)  # 60%宽度
        layout.addLayout(right_panel, 2)  # 20%宽度

        self.setLayout(layout)

    def select_folder(self):
        folder = QFileDialog.getExistingDirectory(self, "选择文件夹")
        if folder:
            self.folder_path.setText(folder)
            # self.load_has_annotator_images(folder)
            self.load_images(folder)



    def load_images(self, folder):
        self.image_list.clear()
        self.image_files = []

        start_num = self.start_spin.value() - 1  # 转换为0-based索引

        # 遍历文件夹
        for root, dirs, files in os.walk(folder):
            image_files = [f for f in files if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
            if not image_files:
                continue

            # 从指定序号开始处理图片
            for i in range(start_num, len(image_files)):
                img_file = image_files[i]
                img_path = os.path.join(root, img_file)
                rel_path = os.path.relpath(img_path, folder)
                self.image_files.append(img_path)
                self.image_list.addItem(rel_path)
                # self.highlight_has_annotator_image_item_color(i)



        if self.image_files:
            self.current_image_index = 0
            self.load_image(0)
            self.output_text.append(f"已加载文件夹: {folder}")
            self.output_text.append(f"共找到 {len(self.image_files)} 张图片")

    def load_image(self, index):
        if 0 <= index < len(self.image_files):
            img_path = self.image_files[index]
            if self.viewer.set_image(img_path):
                self.output_text.append(f"已加载: {img_path}")
                self.current_image_index = index
                self.image_list.setCurrentRow(index)
            else:
                self.output_text.append(f"无法加载: {img_path}")

    def loadPreviousImage(self):
        if self.current_image_index > 0:
            self.loadImageByIndex(self.current_image_index - 1)

            self.viewer.set_annotating(False)
            self.viewer.annotation_mode = None
            self.start_annotate_btn.setText("开始标绘")
            self.output_text.append("结束标绘模式")

            # 重置所有按钮状态
            self.rect_btn.setDown(False)
            self.circle_btn.setDown(False)
            self.text_btn.setDown(False)
            self.polygon_btn.setDown(False)

    def loadNextImage(self):
        if self.current_image_index < len(self.image_files) - 1:
            self.loadImageByIndex(self.current_image_index + 1)

            self.viewer.set_annotating(False)
            self.viewer.annotation_mode = None
            self.start_annotate_btn.setText("开始标绘")
            self.output_text.append("结束标绘模式")

            # 重置所有按钮状态
            self.rect_btn.setDown(False)
            self.circle_btn.setDown(False)
            self.text_btn.setDown(False)
            self.polygon_btn.setDown(False)

    def loadImageByIndex(self, index):
        if 0 <= index < len(self.image_files):
            self.current_image_index = index
            image_path = self.image_files[index]

            # 高亮当前选中的图片
            for i in range(self.image_list.count()):
                if i == index:
                    self.image_list.item(i).setSelected(True)
                else:
                    self.image_list.item(i).setSelected(False)

            try:
                # 清除之前的标注
                self.viewer.clear_annotations()
                self.annotations = []

                self.load_image(index)
            except Exception as e:
                self.output.append(f"加载图片时出错: {str(e)}")

    def on_image_selected(self, item):
        index = self.image_list.row(item)
        self.load_image(index)

    # 加载已经标绘过的图片
    def load_has_annotator_images(self, folder):

        _current_dir = os.path.join(os.path.dirname(folder), '缺陷图片')
        self.has_annotator_files = []
        # 遍历文件夹
        for root, dirs, files in os.walk(_current_dir):
            image_files = [f for f in files if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
            if not image_files:
                continue

            # 从指定序号开始处理图片
            for i in range(0, len(image_files)):
                img_file = image_files[i]
                self.has_annotator_files.append(img_file)


    def highlight_has_annotator_image_item_color(self, currentImageIndex):

        _image_item = self.image_list.item(currentImageIndex)
        _image_item.setForeground(QColor(0, 255, 0))  # 白色文字（如果背景是红色）

        # _c_path = self.image_files[currentImageIndex]
        # _current_file_name = os.path.basename(_c_path)
        # for j in range(0, len(self.has_annotator_files)):
        #     _ann_file = self.has_annotator_files[j]
        #     if os.path.splitext(_current_file_name)[0] in _ann_file:
        #         _image_item = self.image_list.item(currentImageIndex)
        #         _image_item.setForeground(QColor(0, 255, 0))  # 白色文字（如果背景是红色）

    def set_annotation_mode(self, mode):
        if not self.viewer.is_annotating:
            self.output_text.append("请先点击'开始标绘'按钮进入标绘模式")
            return

        if self.viewer.annotation_mode == mode:
            self.viewer.annotation_mode = None
        else:
            self.viewer.annotation_mode = mode

        # 更新按钮状态
        self.rect_btn.setDown(self.viewer.annotation_mode == 'rectangle')
        self.circle_btn.setDown(self.viewer.annotation_mode == 'circle')
        self.text_btn.setDown(self.viewer.annotation_mode == 'text')
        self.polygon_btn.setDown(self.viewer.annotation_mode == 'polygon')

        self.output_text.append(f"当前标注模式: {self.viewer.annotation_mode if self.viewer.annotation_mode else '无'}")

    def toggle_annotation(self):
        if self.viewer.is_annotating:
            self.viewer.set_annotating(False)
            self.viewer.annotation_mode = None
            self.start_annotate_btn.setText("开始标绘")
            self.output_text.append("结束标绘模式")

            # 重置所有按钮状态
            self.rect_btn.setDown(False)
            self.circle_btn.setDown(False)
            self.text_btn.setDown(False)
            self.polygon_btn.setDown(False)
        else:
            self.viewer.set_annotating(True)
            self.start_annotate_btn.setText("结束标绘")
            self.output_text.append("开始标绘模式 - 请选择标注类型")
            self.set_annotation_mode('circle')

    def clear_annotations(self):
        self.viewer.clear_annotations()
        self.output_text.append("已清除所有标注")

    def save_annotations(self):
        if self.current_image_index == -1 or not self.image_files:
            self.output_text.append("没有可保存的图片")
            return

        img_path = self.image_files[self.current_image_index]



        # 弹出对话框获取保存信息
        dialog = AnnotationDialog(self)
        if dialog.exec_() == QDialog.Accepted:
            info = dialog.get_info()

            # 保存标注
            annotated_path, cropped_path = self.viewer.save_annotated_image(img_path, info, self.folder_path.text())

            if annotated_path:
                self.output_text.append(f"已保存标注图片: {annotated_path}")

            if cropped_path:
                self.output_text.append(f"已保存裁剪区域: {cropped_path}")

            # 设置标绘图片高亮
            self.load_has_annotator_images(self.folder_path.text())
            self.highlight_has_annotator_image_item_color(self.current_image_index)

            if not annotated_path and not cropped_path:
                self.output_text.append("没有标注可保存")

    def set_style(self):
        dialog = ColorWidthDialog(self)
        if dialog.exec_() == QDialog.Accepted:
            color, width = dialog.get_settings()
            self.viewer.set_pen_style(color, width)
            self.output_text.append(f"设置标注样式: 颜色 {color.name()}, 线宽 {width}")


